import torch
import torch.nn as nn

from torch.nn import init
from utils.LogUtil import my_logger


# https://github.com/pykao/Modified-3D-UNet-Pytorch
class Modified3DUNet(nn.Module):
    def __init__(self, in_channels, n_classes, base_n_filter=8, if_activate=True):
        super(Modified3DUNet, self).__init__()
        self.in_channels = in_channels
        self.if_activate = if_activate
        self.n_classes = n_classes
        self.base_n_filter = base_n_filter

        self.lrelu = nn.LeakyReLU()
        self.dropout3d = nn.Dropout3d(p=0.6)
        self.upsacle = nn.Upsample(scale_factor=2, mode='nearest')
        self.softMax = nn.Softmax(dim=1)
        self.sigmoid = nn.Sigmoid()

        # Level 1 context pathway
        self.conv3d_c1_1 = nn.Conv3d(self.in_channels, self.base_n_filter, kernel_size=3, stride=1, padding=1,
                                     bias=False)
        self.conv3d_c1_2 = nn.Conv3d(self.base_n_filter, self.base_n_filter, kernel_size=3, stride=1, padding=1,
                                     bias=False)
        self.lrelu_conv_c1 = self.lrelu_conv(self.base_n_filter, self.base_n_filter)
        self.inorm3d_c1 = nn.InstanceNorm3d(self.base_n_filter)

        # Level 2 context pathway
        self.conv3d_c2 = nn.Conv3d(self.base_n_filter, self.base_n_filter * 2, kernel_size=3, stride=2, padding=1,
                                   bias=False)
        self.norm_lrelu_conv_c2 = self.norm_lrelu_conv(self.base_n_filter * 2, self.base_n_filter * 2)
        self.inorm3d_c2 = nn.InstanceNorm3d(self.base_n_filter * 2)

        # Level 3 context pathway
        self.conv3d_c3 = nn.Conv3d(self.base_n_filter * 2, self.base_n_filter * 4, kernel_size=3, stride=2, padding=1,
                                   bias=False)
        self.norm_lrelu_conv_c3 = self.norm_lrelu_conv(self.base_n_filter * 4, self.base_n_filter * 4)
        self.inorm3d_c3 = nn.InstanceNorm3d(self.base_n_filter * 4)

        # Level 4 context pathway
        self.conv3d_c4 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 8, kernel_size=3, stride=2, padding=1,
                                   bias=False)
        self.norm_lrelu_conv_c4 = self.norm_lrelu_conv(self.base_n_filter * 8, self.base_n_filter * 8)
        self.inorm3d_c4 = nn.InstanceNorm3d(self.base_n_filter * 8)

        # Level 5 context pathway, level 0 localization pathway
        self.conv3d_c5 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 16, kernel_size=3, stride=2, padding=1,
                                   bias=False)
        self.norm_lrelu_conv_c5 = self.norm_lrelu_conv(self.base_n_filter * 16, self.base_n_filter * 16)
        self.norm_lrelu_upscale_conv_norm_lrelu_l0 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 16,
                                                                                             self.base_n_filter * 8)

        self.conv3d_l0 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
                                   bias=False)
        self.inorm3d_l0 = nn.InstanceNorm3d(self.base_n_filter * 8)

        # Level 1 localization pathway
        self.conv_norm_lrelu_l1 = self.conv_norm_lrelu(self.base_n_filter * 16, self.base_n_filter * 16)
        self.conv3d_l1 = nn.Conv3d(self.base_n_filter * 16, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
                                   bias=False)
        self.norm_lrelu_upscale_conv_norm_lrelu_l1 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 8,
                                                                                             self.base_n_filter * 4)

        # Level 2 localization pathway
        self.conv_norm_lrelu_l2 = self.conv_norm_lrelu(self.base_n_filter * 8, self.base_n_filter * 8)
        self.conv3d_l2 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 4, kernel_size=1, stride=1, padding=0,
                                   bias=False)
        self.norm_lrelu_upscale_conv_norm_lrelu_l2 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 4,
                                                                                             self.base_n_filter * 2)

        # Level 3 localization pathway
        self.conv_norm_lrelu_l3 = self.conv_norm_lrelu(self.base_n_filter * 4, self.base_n_filter * 4)
        self.conv3d_l3 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 2, kernel_size=1, stride=1, padding=0,
                                   bias=False)
        self.norm_lrelu_upscale_conv_norm_lrelu_l3 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 2,
                                                                                             self.base_n_filter)

        # Level 4 localization pathway
        self.conv_norm_lrelu_l4 = self.conv_norm_lrelu(self.base_n_filter * 2, self.base_n_filter * 2)
        self.conv3d_l4 = nn.Conv3d(self.base_n_filter * 2, self.n_classes, kernel_size=1, stride=1, padding=0,
                                   bias=False)

        self.ds2_1x1_conv3d = nn.Conv3d(self.base_n_filter * 8, self.n_classes, kernel_size=1, stride=1, padding=0,
                                        bias=False)
        self.ds3_1x1_conv3d = nn.Conv3d(self.base_n_filter * 4, self.n_classes, kernel_size=1, stride=1, padding=0,
                                        bias=False)

        # initialize parameters
        self._init_parameters()

    def _init_parameters(self):
        for m in self.modules():
            # print(m)
            if isinstance(m, nn.Conv2d):
                init.xavier_uniform(m.weight)  # uniform distribution
                if m.bias:
                    init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant(m.weight, 1)
                init.constant(m.bias, 0)

    def conv_norm_lrelu(self, feat_in, feat_out):
        return nn.Sequential(
            nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm3d(feat_out),
            nn.LeakyReLU())

    def norm_lrelu_conv(self, feat_in, feat_out):
        return nn.Sequential(
            nn.InstanceNorm3d(feat_in),
            nn.LeakyReLU(),
            nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False))

    def lrelu_conv(self, feat_in, feat_out):
        return nn.Sequential(
            nn.LeakyReLU(),
            nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False))

    def norm_lrelu_upscale_conv_norm_lrelu(self, feat_in, feat_out):
        return nn.Sequential(
            nn.InstanceNorm3d(feat_in),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode='nearest'),
            # should be feat_in*2 or feat_in
            nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm3d(feat_out),
            nn.LeakyReLU())

    def forward(self, x):
        #  Level 1 context pathway
        out = self.conv3d_c1_1(x)  # 3d convolution stride=1
        residual_1 = out
        out = self.lrelu(out)  # leaky relu
        out = self.conv3d_c1_2(out)  # 3d convolution  stride=1
        out = self.dropout3d(out)  # dropout 3d
        out = self.lrelu_conv_c1(out)  # leaky relu and 3d convolution  stride=1
        # Element Wise Summation
        out += residual_1  # addition operation
        context_1 = self.lrelu(out)  # leaky relu
        out = self.inorm3d_c1(out)  # instance 3d normalization
        out = self.lrelu(out)  # leaky relu

        # Level 2 context pathway
        # shape start to reduce
        out = self.conv3d_c2(out)  # 3d convolution stride=2
        residual_2 = out
        out = self.norm_lrelu_conv_c2(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out = self.dropout3d(out)  # dropout 3d
        out = self.norm_lrelu_conv_c2(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out += residual_2  # addition operation
        out = self.inorm3d_c2(out)  # instance 3d normalization
        out = self.lrelu(out)  # leaky relu
        context_2 = out

        # Level 3 context pathway
        out = self.conv3d_c3(out)  # 3d convolution stride=2
        residual_3 = out
        out = self.norm_lrelu_conv_c3(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out = self.dropout3d(out)  # dropout 3d
        out = self.norm_lrelu_conv_c3(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out += residual_3  # addition operation
        out = self.inorm3d_c3(out)  # instance 3d normalization
        out = self.lrelu(out)  # leaky relu
        context_3 = out

        # Level 4 context pathway
        out = self.conv3d_c4(out)  # 3d convolution stride=2
        residual_4 = out
        out = self.norm_lrelu_conv_c4(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out = self.dropout3d(out)  # dropout 3d
        out = self.norm_lrelu_conv_c4(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out += residual_4  # addition operation
        out = self.inorm3d_c4(out)  # instance 3d normalization
        out = self.lrelu(out)  # leaky relu
        context_4 = out

        # Level 5
        out = self.conv3d_c5(out)  # 3d convolution stride=2
        residual_5 = out
        out = self.norm_lrelu_conv_c5(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out = self.dropout3d(out)  # dropout 3d
        out = self.norm_lrelu_conv_c5(out)  # instance 3d normalization,leaky relu and 3d convolution stride=1
        out += residual_5  # addition operation
        # instance 3d normalization,leaky relu,upsample,3d convolution stride=1,instance 3d normalization,leaky relu
        out = self.norm_lrelu_upscale_conv_norm_lrelu_l0(out)

        out = self.conv3d_l0(out)  # 3d convolution kernel_size=1 stride=1
        out = self.inorm3d_l0(out)  # instance 3d normalization
        out = self.lrelu(out)  # leaky relu
        # out 8 64 2 2 2   context_4 8 64 1 1 1
        # Level 1 localization pathway
        out = torch.cat([out, context_4], dim=1)  # concat operation
        out = self.conv_norm_lrelu_l1(out)  # 3d convolution stride=1,instance 3d normalization and leaky relu
        out = self.conv3d_l1(out)  # 3d convolution kernel_size=1 stride=1
        # instance 3d normalization,leaky relu,upsample,3d convolution stride=1,instance 3d normalization,leaky relu
        out = self.norm_lrelu_upscale_conv_norm_lrelu_l1(out)

        # Level 2 localization pathway
        out = torch.cat([out, context_3], dim=1)  # concat operation
        out = self.conv_norm_lrelu_l2(out)  # 3d convolution stride=1, instance 3d normalization,leaky relu
        ds2 = out
        out = self.conv3d_l2(out)  # 3d convolution kernel_size=1 stride=1
        # instance 3d normalization,leaky relu,upsample,3d convolution stride=1,instance 3d normalization,leaky relu
        out = self.norm_lrelu_upscale_conv_norm_lrelu_l2(out)

        # Level 3 localization pathway
        out = torch.cat([out, context_2], dim=1)  # concat operation
        out = self.conv_norm_lrelu_l3(out)  # 3d convolution stride=1, instance 3d normalization,leaky relu
        ds3 = out
        out = self.conv3d_l3(out)  # 3d convolution kernel_size=1 stride=1
        # instance 3d normalization,leaky relu,upsample,3d convolution stride=1,instance 3d normalization,leaky relu
        out = self.norm_lrelu_upscale_conv_norm_lrelu_l3(out)

        # Level 4 localization pathway
        out = torch.cat([out, context_1], dim=1)  # concat operation
        out = self.conv_norm_lrelu_l4(out)  # 3d convolution stride=1, instance 3d normalization,leaky relu
        out_pred = self.conv3d_l4(out)  # 3d convolution kernel_size=1 stride=1

        ds2_1x1_conv = self.ds2_1x1_conv3d(ds2)  # 3d convolution kernel_size=1 stride=1
        ds1_ds2_sum_upscale = self.upsacle(ds2_1x1_conv)  # Upsample
        ds3_1x1_conv = self.ds3_1x1_conv3d(ds3)  # 3d convolution kernel_size=1 stride=1
        ds1_ds2_sum_upscale_ds3_sum = ds1_ds2_sum_upscale + ds3_1x1_conv  # addition
        ds1_ds2_sum_upscale_ds3_sum_upscale = self.upsacle(ds1_ds2_sum_upscale_ds3_sum)  # Upsample

        out = out_pred + ds1_ds2_sum_upscale_ds3_sum_upscale  # addition

        # out = out.permute(0, 2, 3, 4, 1).contiguous().view(-1, self.n_classes)
        # out = out.view(-1, self.n_classes)
        # out = self.softmax(out, dim=1)
        if self.if_activate:
            if self.n_classes > 1:
                result = self.softMax(out)
            elif self.n_classes == 1:
                result = self.sigmoid(out)
            else:
                my_logger.error("The number of output channel is wrong")
                result = None
        else:
            result = out
        return result
        # return result, out
        # out = self.sigmoid(out)
        # return out, seg_layer
